import asyncio
import json
import logging
from typing import Dict, Any, List, Optional
import httpx
import websockets
from fastmcp import FastMCP
from collections import deque
import threading
from datetime import datetime

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 初始化FastMCP实例
mcp = FastMCP("LiveStreamBridge")

# 全局变量
websocket_server = None
server_running = False
clients = set()  # 存储所有连接的客户端，限制不超过3个
data_queue = deque()  # 用于存储待处理的数据队列
queue_lock = threading.Lock()  # 队列锁
processing_task = None  # 数据处理任务

# HTTPS目标服务器地址（可修改）
HTTPS_TARGET_URL = "https://example.com/api/live-stream"

def restructure_websocket_data(raw_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    重构WebSocket接收的数据以满足向HTTPS服务器转发的格式
    
    Args:
        raw_data: WebSocket接收的原始数据
        
    Returns:
        重构后的数据格式
    """
    # 根据实际需求调整数据结构
    structured_data = {
        "timestamp": datetime.now().isoformat(),
        "source": "websocket",
        "payload": raw_data,
        "version": "1.0"
    }
    
    # 如果原始数据包含特定字段，可以进行映射
    if "action" in raw_data:
        structured_data["action"] = raw_data["action"]
    
    if "data" in raw_data:
        structured_data["content"] = raw_data["data"]
    
    return structured_data

def parse_https_response(response_data: Dict[str, Any]) -> Dict[str, Any]:
    """
    解析HTTPS服务器返回的数据格式
    
    Args:
        response_data: HTTPS服务器返回的原始数据
        
    Returns:
        解析后的数据格式，用于传送给大模型
    """
    # 根据实际需求调整解析逻辑
    parsed_data = {
        "timestamp": datetime.now().isoformat(),
        "source": "https_server",
        "status": "success",
        "result": response_data
    }
    
    # 如果响应包含特定字段，可以进行映射
    if "code" in response_data:
        parsed_data["code"] = response_data["code"]
    
    if "message" in response_data:
        parsed_data["message"] = response_data["message"]
    
    if "data" in response_data:
        parsed_data["content"] = response_data["data"]
    
    return parsed_data

async def forward_to_https(data: Dict[str, Any]) -> Dict[str, Any]:
    """
    将数据转发到HTTPS服务器并返回响应
    
    Args:
        data: 要转发的数据
        
    Returns:
        来自HTTPS服务器的响应数据
    """
    async with httpx.AsyncClient() as client:
        try:
            # 重构数据格式
            structured_data = restructure_websocket_data(data)
            
            # 发送POST请求到HTTPS服务器
            response = await client.post(
                HTTPS_TARGET_URL,
                json=structured_data,
                timeout=30.0
            )
            response.raise_for_status()
            
            # 解析响应为JSON
            result = response.json()
            logger.info(f"HTTPS转发成功: {result}")
            return result
            
        except Exception as e:
            logger.error(f"HTTPS转发失败: {e}")
            return {"error": str(e), "status": "failed"}

async def process_data_queue():
    """
    处理数据队列中的数据，逐条转发到HTTPS服务器
    """
    global data_queue, queue_lock
    
    while True:
        try:
            # 检查队列是否有数据
            with queue_lock:
                if not data_queue:
                    await asyncio.sleep(0.1)  # 短暂休眠避免过度占用CPU
                    continue
            
            # 获取队列中的第一条数据
            with queue_lock:
                if data_queue:
                    data_item = data_queue.popleft()
                else:
                    continue
            
            # 转发数据到HTTPS服务器
            response_data = await forward_to_https(data_item)
            
            # 解析响应数据
            parsed_response = parse_https_response(response_data)
            
            # 将响应数据发送给大模型（通过MCP工具）
            await mcp.notify("live_stream_response", parsed_response)
            
            # 短暂休眠避免请求过于频繁
            await asyncio.sleep(0.01)
            
        except Exception as e:
            logger.error(f"处理数据队列时出错: {e}")
            await asyncio.sleep(1)  # 出错时休眠更长时间

async def handle_websocket_client(websocket: websockets.WebSocketServerProtocol, path: str):
    """
    处理WebSocket客户端连接
    
    Args:
        websocket: WebSocket连接对象
        path: 请求路径
    """
    global clients, data_queue, queue_lock
    
    # 检查连接数是否超过限制
    if len(clients) >= 3:
        logger.warning("连接数已达上限，拒绝新连接")
        await websocket.send(json.dumps({
            "error": "服务器连接数已达上限（3个），请稍后再试"
        }, ensure_ascii=False))
        await websocket.close(code=1000, reason="连接数超限")
        return
    
    # 添加客户端到连接集合
    clients.add(websocket)
    client_id = id(websocket)
    logger.info(f"新的直播间客户端连接建立，客户端ID: {client_id}，当前连接数: {len(clients)}")
    
    try:
        # 通知客户端连接已建立
        await websocket.send(json.dumps({
            "status": "connected",
            "message": "直播间WebSocket连接已建立",
            "server": "LiveStreamBridge",
            "clientId": client_id
        }, ensure_ascii=False))
        
        # 持续监听客户端消息
        async for message in websocket:
            try:
                # 解析接收到的JSON数据
                raw_data = json.loads(message)
                logger.info(f"接收到直播间客户端数据: {raw_data}")
                
                # 处理数据（可能是单条或批量数据）
                if isinstance(raw_data, list):
                    # 批量数据，逐条添加到队列
                    with queue_lock:
                        for item in raw_data:
                            data_queue.append(item)
                    logger.info(f"已将{len(raw_data)}条数据添加到处理队列")
                else:
                    # 单条数据，添加到队列
                    with queue_lock:
                        data_queue.append(raw_data)
                    logger.info("已将1条数据添加到处理队列")
                
                # 确认接收数据
                await websocket.send(json.dumps({
                    "status": "received",
                    "message": "数据已接收并加入处理队列"
                }, ensure_ascii=False))
                
            except json.JSONDecodeError:
                error_msg = {"error": "无效的JSON格式"}
                await websocket.send(json.dumps(error_msg, ensure_ascii=False))
                
            except Exception as e:
                logger.error(f"处理直播间客户端消息时出错: {e}")
                error_msg = {"error": f"处理消息时发生错误: {str(e)}"}
                await websocket.send(json.dumps(error_msg, ensure_ascii=False))
                
    except websockets.exceptions.ConnectionClosed:
        logger.info(f"直播间客户端连接已关闭，客户端ID: {client_id}")
    except Exception as e:
        logger.error(f"处理直播间客户端连接时出错: {e}")
    finally:
        # 清理客户端连接
        clients.discard(websocket)
        logger.info(f"直播间客户端连接处理完成，客户端ID: {client_id}，当前连接数: {len(clients)}")

@mcp.tool("send_data", description="向直播间WebSocket客户端发送数据")
async def send_data(data: Dict[str, Any]) -> bool:
    """
    向当前连接的直播间WebSocket客户端发送数据
    
    Args:
        data: 要发送的数据
        
    Returns:
        发送是否成功
    """
    global clients
    
    if not clients:
        logger.warning("当前无直播间客户端连接")
        return False
    
    success_count = 0
    # 向所有连接的客户端发送数据
    for client in list(clients):  # 创建副本以防在迭代时修改集合
        try:
            await client.send(json.dumps(data, ensure_ascii=False))
            success_count += 1
        except Exception as e:
            logger.error(f"向客户端发送数据失败: {e}")
            # 移除已断开的连接
            clients.discard(client)
    
    return success_count > 0

@mcp.tool("get_connection_status", description="获取直播间WebSocket连接状态")
def get_connection_status() -> Dict[str, Any]:
    """
    获取当前直播间WebSocket连接状态
    
    Returns:
        包含连接状态信息的字典
    """
    global clients, data_queue
    
    return {
        "connected_clients": len(clients),
        "max_connections": 3,
        "queue_size": len(data_queue),
        "clients_info": [id(client) for client in clients]
    }

@mcp.tool("start_live_stream_server", description="启动直播间服务")
async def start_live_stream_server() -> Dict[str, Any]:
    """
    启动直播间WebSocket服务器
    
    Returns:
        启动结果信息
    """
    global websocket_server, server_running, processing_task
    
    if server_running:
        return {
            "status": "already_running",
            "message": "直播间服务器已在运行中",
            "address": "ws://localhost:4433"
        }
    
    try:
        # 启动WebSocket服务器
        websocket_server = await websockets.serve(
            handle_websocket_client,
            "localhost",
            4433,
            ping_interval=None  # 禁用自动ping以减少干扰
        )
        
        # 启动数据处理任务
        processing_task = asyncio.create_task(process_data_queue())
        
        server_running = True
        logger.info("直播间WebSocket服务器已在 ws://localhost:4433 启动并开始监听...")
        
        return {
            "status": "success",
            "message": "直播间服务器启动成功",
            "address": "ws://localhost:4433"
        }
    except Exception as e:
        logger.error(f"启动直播间服务器失败: {e}")
        return {
            "status": "error",
            "message": f"启动直播间服务器失败: {str(e)}"
        }

@mcp.tool("stop_live_stream_server", description="停止直播间服务")
async def stop_live_stream_server() -> Dict[str, Any]:
    """
    停止直播间WebSocket服务器
    
    Returns:
        停止结果信息
    """
    global websocket_server, server_running, processing_task, clients, data_queue
    
    if not server_running:
        return {
            "status": "not_running",
            "message": "直播间服务器未在运行"
        }
    
    try:
        # 关闭所有客户端连接
        for client in list(clients):
            try:
                await client.close()
            except:
                pass
        
        # 清理连接集合和数据队列
        clients.clear()
        with queue_lock:
            data_queue.clear()
        
        # 关闭WebSocket服务器
        if websocket_server:
            websocket_server.close()
            await websocket_server.wait_closed()
        
        # 取消数据处理任务
        if processing_task and not processing_task.done():
            processing_task.cancel()
            try:
                await processing_task
            except asyncio.CancelledError:
                pass
        
        server_running = False
        logger.info("直播间WebSocket服务器已关闭")
        
        return {
            "status": "success",
            "message": "直播间服务器已停止"
        }
    except Exception as e:
        logger.error(f"停止直播间服务器失败: {e}")
        return {
            "status": "error",
            "message": f"停止直播间服务器失败: {str(e)}"
        }

@mcp.prompt("连接直播间")
async def connect_live_stream():
    """
    当大模型输入"连接直播间"时触发此函数
    启动MCP服务端并开始监听、转发
    """
    logger.info("收到'连接直播间'指令，正在启动直播间服务...")
    
    # 启动直播间服务器
    result = await start_live_stream_server()
    
    if result["status"] == "success":
        # 启动FastMCP服务
        logger.info("直播间服务启动成功，正在启动MCP服务...")
        return {
            "status": "success",
            "message": "直播间连接已建立，MCP服务已启动",
            "server_address": result["address"]
        }
    else:
        return {
            "status": "error",
            "message": f"直播间服务启动失败: {result['message']}"
        }

async def main():
    """
    主函数：启动FastMCP服务
    """
    logger.info("正在启动直播间MCP服务...")
    await mcp.run()

if __name__ == "__main__":
    asyncio.run(main())